#include "analysis.h"
#include "utility.h"
#include <fstream>
#include <iostream>
#include <sstream>

void ReadTrajectory(const std::string &traj_file, std::vector<int> &steps, std::vector<MolSys> &mols)
{
    std::ifstream infile(traj_file, std::ios::in);
    if (!infile)
    {
        std::cerr << "Error: file '" << traj_file << "' does NOT exist" << std::endl;
        exit(1);
    }
    MolSys mol;
    std::string line;
    int nLine = 0;
    steps.clear();
    mols.clear();
    while (std::getline(infile, line))
    {
        if (line[0] == '#') // comment line
        {
            continue;
        }
        if (line == "") // blank line
        {
            continue;
        }
        nLine += 1;
        if (nLine == 1)
        {
           // mol.natoms_ = std::stoi(line);
            std::stringstream ss(line);
            ss >> mol.natoms_;
        }
        //   outfile << "iStep = " << count << ", Ep(au) = " << potential_energy_  << ", Ek(au) = " << kinetic_energy_  << ", Etot(au) = " << total_energy_     << ", T(K) = " << temperature_    << std::endl;
        std::stringstream ss(line);
        if (nLine == 2)
        {
            int iStep = -1;
            std::string dump;
            ss >> dump >> dump >> iStep >> dump >> dump >> dump >> mol.potential_energy_ >> dump >> dump >> dump >> mol.kinetic_energy_ >> dump >> dump >> dump >> mol.total_energy_ >> dump >> dump >> dump >> mol.temperature_;
            mol.atoms_.clear();
            steps.push_back(iStep);
        }
        if (nLine > 2)
        {
            Atom atom;
            ss >> atom.name_ >> atom.coord_[0] >> atom.coord_[1] >> atom.coord_[2] >> atom.veloc_[0] >> atom.veloc_[1] >> atom.veloc_[2];
            mol.atoms_.push_back(atom);
        }
        if (nLine == 2 + mol.natoms_)
        {
            mols.push_back(mol);
            nLine = 0;
        }
    } // loop end: getline(infile, line)
    return;
}
/*
autocorrelation function Cxx(t)=<x(0)x(t)>
@param acf: acf[0]  =  Cxx(0); acf[1] = Cxx(1); ...; acf[t]=Cxx(t)
*/
std::vector<double> AutoCorrelationFunc(const std::vector<double> elem, const int iTtotal)
{
    if (int(elem.size()) < iTtotal)
    {
        std::cerr << "Error: iTtotal should NOT be greater than elem.size()" << std::endl;
        exit(1);
    }
    std::vector<double> acf(iTtotal);
    for (int k = 0; k < iTtotal; k++)
    {
        for (int i = 0, j = k; j < iTtotal; i++, j++)
        {
            acf[k] += elem[i] * elem[j];
        }
        acf[k] /= iTtotal - k;
    }
    return acf;
}
/*
@param dt: time step in simulation
*/
double SelfDiffusionCoefficient(const std::string &traj_file, const double &dt)
{
    std::vector<MolSys> mols;
    std::vector<int> steps;
    ReadTrajectory(traj_file, steps, mols);
    if (steps.size() != mols.size())
    {
        std::cerr << "Error in ReadTrajectory" << std::endl;
        exit(1);
    }
    int nsteps = steps.size();
    int natoms = mols[0].natoms_;
    // for each atom in each mol
    std::vector<std::vector<double>> vxs;
    std::vector<std::vector<double>> vys;
    std::vector<std::vector<double>> vzs;
    for (int k = 0; k < natoms; k++)
    {
        std::vector<double> vx, vy, vz;
        // atom k velocities at steps
        for (int i = 0; i < nsteps; i++)
        {
            vx.push_back(mols[i].atoms_[k].veloc_[0]);
            vy.push_back(mols[i].atoms_[k].veloc_[1]);
            vz.push_back(mols[i].atoms_[k].veloc_[2]);
        }
        vxs.push_back(vx);
        vys.push_back(vy);
        vzs.push_back(vz);
    }
    double time_interval = dt * (steps[1] - steps[0]);
    return SelfDiffusionCoefficient(vxs, vys, vzs, time_interval);
}

/*
@param 
*/
double SelfDiffusionCoefficient(const std::vector<std::vector<double>> &vxs,
                                const std::vector<std::vector<double>> &vys,
                                const std::vector<std::vector<double>> &vzs,
                                const double &time_interval)
{
    int natoms = vxs.size();
    int nsteps = vxs[0].size();
    std::vector<double> avg_vx_acf(nsteps, 0);
    std::vector<double> avg_vy_acf(nsteps, 0);
    std::vector<double> avg_vz_acf(nsteps, 0);
    for (int k = 0; k < natoms; k++)
    {
        std::vector<double> vx_acf = AutoCorrelationFunc(vxs[k], nsteps);
        std::vector<double> vy_acf = AutoCorrelationFunc(vys[k], nsteps);
        std::vector<double> vz_acf = AutoCorrelationFunc(vzs[k], nsteps);
        for (int i = 0; i < nsteps; i++)
        {
            avg_vx_acf[i] += vx_acf[i];
            avg_vy_acf[i] += vy_acf[i];
            avg_vz_acf[i] += vz_acf[i];
        }
    }
    // average velocity autocorrelation functions
    for (int i = 0; i < nsteps; i++)
    {
        avg_vx_acf[i] /= (double)natoms;
        avg_vy_acf[i] /= (double)natoms;
        avg_vz_acf[i] /= (double)natoms;
    }
    // calculate self diffusion coefficient
    double d_x = SelfDiffusionCoefficient(avg_vx_acf, time_interval);
    double d_y = SelfDiffusionCoefficient(avg_vy_acf, time_interval);
    double d_z = SelfDiffusionCoefficient(avg_vz_acf, time_interval);
    double d = (d_x + d_y + d_z) / 3.0;
    return d;
}

/*
D = \int_{0}^{\infty} Cvv(t) dt
D = (Cvv(0) + Cvv(dt) + Cvv(2*dt) + ... + Cvv(n*dt)) * dt
*/
double SelfDiffusionCoefficient(const std::vector<double> &vacf, const double &time_interval)
{
    // integrate
    double d = 0;
    for (size_t i = 0; i < vacf.size(); ++i)
    {
        d += vacf[i];
    }
    d *= time_interval;
    return d;
}
